#include "StdAfx.h"
#include "StratumInterp.h"

//reference: http://www.nbb.cornell.edu/neurobio/land/OldstudentProjects/cs490-94to95/clang/kriging.html
//reference: http://wenku.baidu.com/view/36509512a8114431b90dd861.html

void InitialzeCoeffs(
	const CArray2D<glm::vec3>& cDrillData, 
	UINT uiStratum, 
	float& c0, 
	float& c1, 
	float& a);

typedef enum EVariogramMode
{
	EVariogramModeSphere = 0, 
	EVariogramModeExponent
};

float CalculateVariogram(
	float c0, 
	float c1, 
	float a, 
	EVariogramMode eMode, 
	float h);

void SolveMatrixEquation(
	const CArray2D<float>& v, 
	const float* d, 
	float* w);

HRESULT StratumInterpKriging(
	const CArray2D<glm::vec3>& cDrillData, 
	CArray3D<glm::vec3>& cStratumData)
{
	glm::vec3 stBoundMin, stBoundMax;
	UINT uiMeshWidth = cStratumData.GetWidth() - 1;
	UINT uiMeshHeight = cStratumData.GetHeight() - 1;

	GetDrillBoundBox(cDrillData, stBoundMin, stBoundMax);

	UINT uiDrillCount = cDrillData.GetHeight();
	UINT uiStratumCount = cDrillData.GetWidth();
	float fCellWidth = (stBoundMax.x - stBoundMin.x) / uiMeshWidth;
	float fCellHeight = (stBoundMax.y - stBoundMin.y) / uiMeshHeight;

	for(UINT uiStratum = 0; uiStratum < uiStratumCount; ++uiStratum)
	{
		float c0 = 0, c1 = 0, a = 0;
		InitialzeCoeffs(cDrillData, uiStratum, c0, c1, a);

		//build V matrix
		CArray2D<float> v(uiDrillCount + 1, uiDrillCount + 1);
		for(UINT i = 0; i < uiDrillCount; ++i)
		{
			for(UINT j = i; j < uiDrillCount; ++j)
			{
				float h = (cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) * 
					(cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) + 
					(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y) * 
					(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y);
				h = sqrt(h);
				v[i][j] = v[j][i] = CalculateVariogram(c0, c1, a, EVariogramModeSphere, h);
			}
		}
		for(UINT i = 0; i < uiDrillCount; ++i)
		{
			v[uiDrillCount][i] = 1;  //The components of last row are 1
			v[i][uiDrillCount] = 1;  //The components of last column are 1
		}
		v[uiDrillCount][uiDrillCount] = 0; //the last component of the matrix is 0

		//interpolation
		float* d = new float[uiDrillCount + 1];
		float* w = new float[uiDrillCount + 1];
		for(UINT i = 0; i < uiMeshHeight + 1; ++i)
		{
			for(UINT j = 0; j < uiMeshWidth + 1; ++j)
			{
				cStratumData[uiStratum][i][j].x = stBoundMin.x + j * fCellWidth;
				cStratumData[uiStratum][i][j].y = stBoundMin.y + i * fCellHeight;

				//build D vector
				for(UINT k = 0; k < uiDrillCount; ++k)
				{
					float h = (cDrillData[k][uiStratum].x - cStratumData[uiStratum][i][j].x) * 
						(cDrillData[k][uiStratum].x - cStratumData[uiStratum][i][j].x) + 
						(cDrillData[k][uiStratum].y - cStratumData[uiStratum][i][j].y) * 
						(cDrillData[k][uiStratum].y - cStratumData[uiStratum][i][j].y);
					h = sqrt(h);
					d[k] = CalculateVariogram(c0, c1, a, EVariogramModeSphere, h);
				}
				d[uiDrillCount] = 1; //The last component of the vector is 1.

				//solve matrix equation to get W
				SolveMatrixEquation(v, d, w);

				//get interpolated z
				cStratumData[uiStratum][i][j].z = 0;
				for(UINT k = 0; k < uiDrillCount; ++k)
				{
					cStratumData[uiStratum][i][j].z += w[k] * cDrillData[k][uiStratum].z;
				}
			}
		}
		delete[] d;
		delete[] w;
	}

	return S_OK;
}

typedef struct
{
	float fVarigram;
	float fSumOfWeight;
	int nCount;
}SVarigramData;

void InitialzeCoeffs(
	const CArray2D<glm::vec3>& cDrillData, 
	UINT uiStratum, 
	float& c0, 
	float& c1, 
	float& a)
{
	UINT uiDrillCount = cDrillData.GetHeight();
	
	float halfMaxDist = 0;
	int nt = 25;
	float ht = 0;
	for(UINT i = 0; i < uiDrillCount - 1; ++i)
	{
		for(UINT j = i + 1; j < uiDrillCount; ++j)
		{
			float d = (cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) * 
				(cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) + 
				(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y) * 
				(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y);
			d = sqrt(d);
			halfMaxDist = glm::max(halfMaxDist, d / 2.0f);
		}
	}
	ht = halfMaxDist / nt;

	//calculate known variogram result
	//????why using this method????
	SVarigramData* pVarigramData = new SVarigramData[nt];
	::ZeroMemory(pVarigramData, sizeof(SVarigramData) * nt);
	for(UINT i = 0; i < uiDrillCount - 1; ++i)
	{
		for(UINT j = i + 1; j < uiDrillCount; ++j)
		{
			float d = (cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) * 
				(cDrillData[i][uiStratum].x - cDrillData[j][uiStratum].x) + 
				(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y) * 
				(cDrillData[i][uiStratum].y - cDrillData[j][uiStratum].y);
			d = sqrt(d);

			if(d > __EPSILON__ && d <= halfMaxDist)
			{
				int hp = (int)round(d / ht);
				float fWeight = 1 - (d - hp * ht) * (d - hp * ht) / ht / ht / 4;
				if(hp > 0)
				{
					pVarigramData[hp - 1].fVarigram += 
						(cDrillData[i][uiStratum].z - cDrillData[j][uiStratum].z) * 
						(cDrillData[i][uiStratum].z - cDrillData[j][uiStratum].z) * fWeight;
					pVarigramData[hp - 1].fSumOfWeight += fWeight; 
					++pVarigramData[hp - 1].nCount;
				}
			}
		}
	}
	for(int i = 0; i < nt; ++i)
	{
		if(pVarigramData[i].fSumOfWeight != 0)
		{
			pVarigramData[i].fVarigram /= pVarigramData[i].fSumOfWeight * 2;
		}
	}

	//fitting variogram equation(sphere mode)
	//r = c0 + c1 * (1.5f * h / a + - 0.5f * (h / a) ^ 3)
	//r = b0 + b1 * h + b2 * h^3
	//c0 = b0
	//c1 = b1 / 3.0f * sqrt(b1 / b2 / - 3.0f)
	//a = sqrt(b1 / b2 / - 3.0f)
	//according to following equations, to solve b1 and b2
	//C(h, h) * b1 + C(h, h^3) * b2 = C(h, r)
	//C(h^3, h) * b1 + C(h^3, h^3) * b2 = C(h^3, r)
	//according to following equation, to solve b0
	//b0 = E(r) - b1 * E(h) - b2 * E(h^3)
	int nTotalCount = 0;
	float E_r = 0;
	float E_h = 0;
	float E_h3 = 0;
	float E_h_h = 0;
	float E_h_h3 = 0;
	float E_h3_h3 = 0;
	float E_h_r = 0;
	float E_h3_r = 0;

	for(int i = 0; i < nt; ++i)
	{
		float h = (i + 1) * ht;
		float h3 = h * h * h; 

		nTotalCount += pVarigramData[i].nCount;
		E_r += pVarigramData[i].nCount * pVarigramData[i].fVarigram;
		E_h += pVarigramData[i].nCount * h;
		E_h3 += pVarigramData[i].nCount * h3;
		E_h_h += pVarigramData[i].nCount * h * h;
		E_h_h3 += pVarigramData[i].nCount * h * h3;
		E_h3_h3 += pVarigramData[i].nCount * h3 * h3;
		E_h_r += pVarigramData[i].nCount * h * pVarigramData[i].fVarigram;
		E_h3_r += pVarigramData[i].nCount * h3 * pVarigramData[i].fVarigram;
	}
	E_r /= nTotalCount;
	E_h /= nTotalCount;
	E_h3 /= nTotalCount;

	float C_h_h = E_h_h - E_h * E_h / nTotalCount;
	float C_h_h3 = E_h_h3 - E_h * E_h3 / nTotalCount;
	float C_h3_h3 = E_h3_h3 - E_h3 * E_h3 / nTotalCount;
	float C_h_r = E_h_r - E_h * E_r / nTotalCount;
	float C_h3_r = E_h3_r - E_h3 * E_r / nTotalCount;
	float b0 = 0, b1 = 0, b2 = 0;

	b1 = (C_h_r * C_h3_h3 - C_h3_r * C_h_h3) / (C_h_h * C_h3_h3 - C_h_h3 * C_h_h3);
	b2 = (C_h3_r * C_h_h - C_h_r * C_h_h3) / (C_h_h * C_h3_h3 - C_h_h3 * C_h_h3);
	b0 = E_r - b1 * E_h - b2 * E_h3;
	if(b0 < 0)
	{
		C_h_h = E_h_h;
		C_h_h3 = E_h_h3;
		C_h3_h3 = E_h3_h3;
		C_h_r = E_h_r;
		C_h3_r = E_h3_r;
		b1 = (C_h_r * C_h3_h3 - C_h3_r * C_h_h3) / (C_h_h * C_h3_h3 - C_h_h3 * C_h_h3);
		b2 = (C_h3_r * C_h_h - C_h_r * C_h_h3) / (C_h_h * C_h3_h3 - C_h_h3 * C_h_h3);
		b0 = 0;
	}

	//output parameters
	c0 = b0;
	c1 = b1 * sqrt(-b1 / b2 / 3.0f) / 3.0f;
	a = sqrt(-b1 / b2 / 3.0f); //here comparing with origin impelementation, dividing by 2 is removed.

	delete[] pVarigramData;
}

float CalculateVariogram(
	float c0, 
	float c1, 
	float a, 
	EVariogramMode eMode, 
	float h)
{
	float r = 0;

	switch(eMode)
	{
	case EVariogramModeSphere:
		/*
			r = c0 + c1 * (1.5f * h / a + - 0.5f * (h / a) ^ 3)
			calcuating method is weired???
		*/
		if(h < __EPSILON__)
		{
			r = c0 + c1;
		}
		else if(h <= a)
		{
			r = c1 - c1 * (1.5f * h / a - 0.5f * h * h * h / a / a / a);
		}
		else
		{
			r = 0;
		}
		break;
	case EVariogramModeExponent:
		/*
			r = c0 + c1 * (1 - exp(-3 * |h| / a)
		*/
		r = c0 + c1 * (1 - exp(-3 * abs(h) / a));
		break;
	default:
		assert(FALSE);
	}

	return r;
}

void SolveMatrixEquation(
	const CArray2D<float>& v, 
	const float* d, 
	float* w)
{
	int n = (int)v.GetWidth();
	float* y = new float[n];
	CArray2D<float> l(n, n);
	CArray2D<float> u(n, n);

	for(int i = 0; i < n; ++i)
	{
		u[i][i] = 1;
	}

	for(int i = 0; i < n; ++i)
	{
		for(int j = i; j < n; ++j)
		{
			l[j][i] = v[j][i];
			for(int k = 0; k < i; ++k)
			{
				l[j][i] -= l[j][k] * u[k][i];
			}
		}
		for(int j = i + 1; j < n; ++j)
		{
			u[i][j] = v[i][j];
			for(int k = 0; k < i; ++k)
			{
				u[i][j] -= l[i][k] * u[k][j];
			}
			u[i][j] /= l[i][i];
		}
	}

	for(int i = 0; i < n; ++i)
	{
		y[i] = d[i];
		for(int j = 0; j < i; ++j)
		{
			y[i] -= l[i][j] * y[j];
		}
		y[i] /= l[i][i];
	}

	for(int i = n - 1; i >= 0; --i)
	{
		w[i] = y[i];
		for(int j = i + 1; j < n; ++j)
		{
			w[i] -= u[i][j] * w[j];
		}
	}

	delete[] y;
}